import os

# 导入核心模块
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict, Optional
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
# 加载环境变量（包含API密钥）
load_dotenv()



# 1. 定义状态结构：使用TypedDict确保数据类型安全
class InputState(TypedDict):
    question: str  # 接收用户问题
    llm_answer: Optional[str]  # 存储LLM生成的中文回答，可选（初始为None）


class OutputState(TypedDict):
    answer: str  # 最终输出的英文翻译结果


# 合并输入与输出状态，形成全局状态
class OverallState(InputState, OutputState):
    pass


# 2. 定义节点函数：每个函数接收状态，处理后返回更新的状态
def llm_node(state: InputState):
    """LLM节点：接收用户问题，生成中文回答"""
    # 构建对话消息
    messages = [
        ("system", "你是一位乐于助人的智能小助理，用中文回答用户问题，回答需简洁准确。"),
        ("human", state["question"])
    ]
    # 初始化DeepSeek模型（可替换为其他模型）
    llm = ChatOpenAI(
        model= os.getenv("model"),
        base_url= os.getenv("dashscope_base_url"),
        temperature=0, # 温度设为0，确保回答稳定
        api_key = os.getenv("dashscope_api_key")
    )
    # 调用LLM生成回答
    response = llm.invoke(messages)
    # 将结果写入状态，返回更新后的状态
    return {"llm_answer": response.content}


def translate_node(state: InputState):
    """翻译节点：将LLM生成的中文回答翻译成英文"""
    messages = [
        ("system", "无论接收的文本是什么语言，都请准确翻译成英文，保持语义完整。"),
        ("human", state["llm_answer"])
    ]
    # 复用LLM实例（实际开发中可优化为全局实例，减少初始化开销）
    llm = ChatOpenAI(
        model= os.getenv("model"),
        base_url= os.getenv("dashscope_base_url"),
        temperature=0,
        api_key = os.getenv("dashscope_api_key")
    )
    response = llm.invoke(messages)
    # 将英文翻译结果写入状态的output字段
    return {"answer": response.content}


# 3. 构建图结构
# 初始化StateGraph，指定全局状态、输入与输出模式
builder = StateGraph(
    state_schema=OverallState,
    input_schema=InputState,
    output_schema=OutputState
)

# 添加节点：指定节点名称与对应的函数
builder.add_node("llm_node", llm_node)  # LLM生成中文回答节点
builder.add_node("translate_node", translate_node)  # 翻译节点

# 4. 连接边：定义执行流程
builder.add_edge(START, "llm_node")  # 入口点→LLM节点
builder.add_edge("llm_node", "translate_node")  # LLM节点→翻译节点
builder.add_edge("translate_node", END)  # 翻译节点→结束点

# 5. 编译图：验证结构合法性，生成可执行实例
graph = builder.compile()

# 6. 调用图：传入用户问题，获取结果
user_question = "请简单介绍LangGraph的核心功能"
result = graph.invoke({"question": user_question})

# 输出最终结果
print("用户问题：", user_question)
print("英文回答：", result["answer"])